import os
import time
import pdfplumber
import re
from datetime import datetime
from typing import List, Tuple, Optional

def detect_columns(page, min_column_width=200):
    """
    检测页面是否为分栏格式并返回栏的边界
    
    Args:
        page: pdfplumber页面对象
        min_column_width: 最小栏宽度（像素）
    
    Returns:
        list: 栏的边界列表 [(x0, x1), ...] 或 None（如果不是分栏）
    """
    # 获取页面尺寸
    page_width = page.width
    page_height = page.height
    
    # 提取所有文本对象的位置信息
    chars = page.chars
    if not chars:
        return None
    
    # 按x坐标分组文本
    x_positions = [char['x0'] for char in chars]
    x_positions.sort()
    
    # 寻找明显的间隙来识别栏分界
    gaps = []
    for i in range(1, len(x_positions)):
        gap = x_positions[i] - x_positions[i-1]
        if gap > 20:  # 间隙大于20像素认为可能是栏分界
            gaps.append((x_positions[i-1], x_positions[i], gap))
    
    # 如果没有明显间隙，尝试按页面宽度平分
    if not gaps:
        # 检查是否可能是双栏
        left_chars = [char for char in chars if char['x0'] < page_width / 2]
        right_chars = [char for char in chars if char['x0'] >= page_width / 2]
        
        if len(left_chars) > 10 and len(right_chars) > 10:
            return [(0, page_width / 2), (page_width / 2, page_width)]
    
    # 分析间隙，确定栏边界
    columns = []
    if gaps:
        # 找到最大的间隙作为主要分栏分界
        main_gap = max(gaps, key=lambda x: x[2])
        if main_gap[2] > 30:  # 间隙足够大
            columns = [(0, main_gap[0]), (main_gap[1], page_width)]
    
    # 验证栏宽度
    valid_columns = []
    for col in columns:
        if col[1] - col[0] >= min_column_width:
            valid_columns.append(col)
    
    return valid_columns if len(valid_columns) >= 2 else None

def extract_column_text(page, column_bounds):
    """
    从指定栏边界提取文本
    
    Args:
        page: pdfplumber页面对象
        column_bounds: 栏边界 (x0, x1)
    
    Returns:
        str: 提取的文本
    """
    x0, x1 = column_bounds
    
    # 创建栏区域的裁剪
    column_area = page.crop((x0, 0, x1, page.height))
    
    # 提取文本
    text = column_area.extract_text()
    return text if text else ""

def extract_text_with_layout(page):
    """
    智能提取页面文本，自动处理分栏格式
    
    Args:
        page: pdfplumber页面对象
    
    Returns:
        str: 按正确顺序提取的文本
    """
    # 检测是否为分栏格式
    columns = detect_columns(page)
    
    if columns:
        # 分栏格式：按栏顺序提取文本
        column_texts = []
        for column_bounds in columns:
            column_text = extract_column_text(page, column_bounds)
            if column_text.strip():
                column_texts.append(column_text.strip())
        
        # 将各栏文本合并，栏之间用双换行分隔
        return "\n\n--- 栏分隔 ---\n\n".join(column_texts)
    else:
        # 单栏格式：直接提取
        return page.extract_text() or ""

def clean_text_advanced(text):
    """
    高级文本清理，专门处理分栏PDF的常见问题
    
    Args:
        text: 原始文本
    
    Returns:
        str: 清理后的文本
    """
    if not text:
        return ""
    
    # 移除页眉页脚（通常是重复的短行）
    lines = text.split('\n')
    cleaned_lines = []
    
    for line in lines:
        line = line.strip()
        if not line:
            continue
            
        # 跳过可能的页码
        if re.match(r'^\d+$', line) and len(line) <= 3:
            continue
            
        # 跳过过短的行（可能是页眉页脚）
        if len(line) < 10 and not re.search(r'[。！？]', line):
            continue
            
        cleaned_lines.append(line)
    
    # 重新组合文本
    text = '\n'.join(cleaned_lines)
    
    # 修复分栏导致的断句问题
    # 如果行末没有标点符号且下一行不是以大写字母或数字开头，则连接
    lines = text.split('\n')
    merged_lines = []
    
    i = 0
    while i < len(lines):
        current_line = lines[i].strip()
        
        if i < len(lines) - 1:
            next_line = lines[i + 1].strip()
            
            # 检查是否需要合并行
            if (current_line and next_line and 
                not current_line.endswith(('。', '！', '？', '.', '!', '?', ':', ';')) and
                not next_line[0].isupper() and
                not next_line[0].isdigit() and
                '--- 栏分隔 ---' not in current_line and
                '--- 栏分隔 ---' not in next_line):
                
                merged_lines.append(current_line + ' ' + next_line)
                i += 2  # 跳过下一行
                continue
        
        merged_lines.append(current_line)
        i += 1
    
    # 最终清理
    text = '\n'.join(merged_lines)
    text = re.sub(r'\n\s*\n\s*\n', '\n\n', text)  # 移除多余空行
    text = re.sub(r' +', ' ', text)  # 移除多余空格
    
    return text.strip()

def pdf_to_markdown_column_aware(pdf_path, output_path=None):
    """
    将PDF文件转换为Markdown格式，智能处理分栏布局
    
    Args:
        pdf_path: PDF文件路径
        output_path: 输出Markdown文件路径
    
    Returns:
        bool: 转换是否成功
    """
    if not os.path.exists(pdf_path):
        print(f"错误：文件 {pdf_path} 不存在")
        return False
    
    if output_path is None:
        base_name = os.path.splitext(pdf_path)[0]
        output_path = f"{base_name}_column_aware.md"
    
    try:
        with pdfplumber.open(pdf_path) as pdf:
            markdown_content = []
            total_pages = len(pdf.pages)
            
            print(f"开始处理 {pdf_path}，共 {total_pages} 页...")
            
            for i, page in enumerate(pdf.pages, 1):
                print(f"处理第 {i}/{total_pages} 页...")
                
                # 使用智能布局提取
                text = extract_text_with_layout(page)
                
                if text:
                    # 高级文本清理
                    cleaned_text = clean_text_advanced(text)
                    if cleaned_text:
                        # 添加页面分隔符
                        page_content = f"\n\n--- 第 {i} 页 ---\n\n{cleaned_text}"
                        markdown_content.append(page_content)
            
            # 写入文件
            with open(output_path, 'w', encoding='utf-8') as f:
                f.write('\n\n'.join(markdown_content))
            
            print(f"转换完成：{output_path}")
            return True
    
    except Exception as e:
        print(f"转换 {pdf_path} 时出现错误：{str(e)}")
        return False

def batch_convert_column_aware(folder_path, output_folder=None):
    """
    批量转换PDF文件，智能处理分栏格式
    
    Args:
        folder_path: 输入文件夹路径
        output_folder: 输出文件夹路径
    
    Returns:
        dict: 转换统计信息
    """
    if not os.path.exists(folder_path):
        print(f"错误：文件夹 {folder_path} 不存在")
        return {}
    
    # 查找PDF文件
    pdf_files = []
    for root, dirs, files in os.walk(folder_path):
        for file in files:
            if file.lower().endswith('.pdf'):
                pdf_files.append(os.path.join(root, file))
    
    if not pdf_files:
        print(f"在 {folder_path} 中未找到PDF文件")
        return {}
    
    # 设置输出文件夹
    if output_folder is None:
        output_folder = os.path.join(folder_path, "markdown_output_column_aware")
    
    os.makedirs(output_folder, exist_ok=True)
    
    # 转换统计
    stats = {
        'total': len(pdf_files),
        'success': 0,
        'failed': 0,
        'start_time': time.time()
    }
    
    print(f"找到 {len(pdf_files)} 个PDF文件，开始批量转换...")
    
    for i, pdf_path in enumerate(pdf_files, 1):
        print(f"\n[{i}/{len(pdf_files)}] 处理: {os.path.basename(pdf_path)}")
        
        # 生成输出路径
        base_name = os.path.splitext(os.path.basename(pdf_path))[0]
        output_path = os.path.join(output_folder, f"{base_name}_column_aware.md")
        
        # 转换文件
        if pdf_to_markdown_column_aware(pdf_path, output_path):
            stats['success'] += 1
        else:
            stats['failed'] += 1
    
    # 计算耗时
    stats['duration'] = time.time() - stats['start_time']
    
    # 输出统计信息
    print(f"\n=== 转换完成 ===")
    print(f"总文件数: {stats['total']}")
    print(f"成功转换: {stats['success']}")
    print(f"转换失败: {stats['failed']}")
    print(f"总耗时: {stats['duration']:.2f} 秒")
    print(f"输出目录: {output_folder}")
    
    return stats

def main():
    """
    主函数
    """
    # 设置输入和输出文件夹
    input_folder = r"C:\Users\ajunfornet\Desktop\PubMed_556"  # 修改为你的PDF文件夹路径
    output_folder = r"C:\Users\ajunfornet\Desktop\PubMed_556\markdown_output_column_aware"  # 修改为你的输出文件夹路径
    
    print("=== 分栏感知PDF转Markdown工具 ===")
    print(f"输入文件夹: {input_folder}")
    print(f"输出文件夹: {output_folder}")
    print("\n特性:")
    print("- 自动检测分栏格式")
    print("- 智能文本提取和排序")
    print("- 高级文本清理")
    print("- 保持原文档逻辑结构")
    
    # 检查输入文件夹
    if not os.path.exists(input_folder):
        print(f"\n警告：输入文件夹 {input_folder} 不存在")
        print("请修改 main() 函数中的 input_folder 路径")
        return
    
    # 开始转换
    print("\n开始转换...")
    stats = batch_convert_column_aware(input_folder, output_folder)
    
    if stats:
        print("\n转换完成！")
    else:
        print("\n转换失败或无文件可转换")

if __name__ == "__main__":
    main()